import heapq
import time
from collections import deque
from enum import Enum
from typing import List, Tuple, NamedTuple, Callable, Optional, Generic, TypeVar, Set, Deque, Dict
import random
from math import sqrt


class Function:
    def __init__(self, func):
        self.func = func

    def __call__(self, *args, **kwargs):
        start_time = time.time()
        self.ret = self.func(*args, **kwargs)
        self.spend_time = time.time() - start_time
        return self


class Cell(str, Enum):
    EMPTY = " "
    BLOCKED = "O"
    START = "S"
    GOAL = "G"
    PATH = "*"


class MazeLocation(NamedTuple):
    row: int
    column: int


class Maze:
    def __init__(self, rows: int = 20, columns: int = 20, sparseness: float = 0.2,
                 start: MazeLocation = MazeLocation(0, 0), goal: MazeLocation = MazeLocation(19, 19)):
        self._rows: int = rows
        self._columns: int = columns
        self.start: MazeLocation = start
        self.goal = goal
        self._grid = [[Cell.EMPTY for c in range(columns)] for c in range(rows)]
        self._randomly_fill(rows, columns, sparseness)
        self._grid[start.row][start.column] = Cell.START
        self._grid[goal.row][goal.column] = Cell.GOAL

    def _randomly_fill(self, rows: int, columns: int, sparseness: float):
        for row in range(rows):
            for column in range(columns):
                if random.uniform(0, 1.0) < sparseness:
                    self._grid[row][column] = Cell.BLOCKED

    def __str__(self):
        output: str = ""
        for row in self._grid:
            output += "".join([c.value for c in row]) + "\n"
        return output

    def goal_test(self, ml: MazeLocation) -> bool:
        return ml == self.goal

    def successors(self, ml: MazeLocation) -> List[MazeLocation]:
        locations: List[MazeLocation] = []
        if ml.row + 1 < self._rows and self._grid[ml.row + 1][ml.column] != Cell.BLOCKED:
            locations.append(MazeLocation(ml.row + 1, ml.column))
        if ml.row - 1 > 0 and self._grid[ml.row - 1][ml.column] != Cell.BLOCKED:
            locations.append(MazeLocation(ml.row - 1, ml.column))
        if ml.column + 1 < self._columns and self._grid[ml.row][ml.column + 1] != Cell.BLOCKED:
            locations.append(MazeLocation(ml.row, ml.column + 1))
        if ml.column - 1 > 0 and self._grid[ml.row][ml.column - 1] != Cell.BLOCKED:
            locations.append(MazeLocation(ml.row, ml.column - 1))
        return locations

    def mark(self, path: List[MazeLocation]):
        for maze_location in path:
            self._grid[maze_location.row][maze_location.column] = Cell.PATH
        self._grid[self.start.row][self.start.column] = Cell.START
        self._grid[self.goal.row][self.goal.column] = Cell.GOAL

    def clear(self, path: List[MazeLocation]):
        for maze_location in path:
            self._grid[maze_location.row][maze_location.column] = Cell.EMPTY
        self._grid[self.start.row][self.start.column] = Cell.START
        self._grid[self.goal.row][self.goal.column] = Cell.GOAL


T = TypeVar('T')


class Stack(Generic[T]):
    def __init__(self):
        self._container: List[T] = []

    @property
    def empty(self) -> bool:
        return not self._container

    def push(self, item: T):
        self._container.append(item)

    def pop(self):
        return self._container.pop()


class Node(Generic[T]):
    def __init__(self, state: T, parent: Optional['Node'], cost: float = 0.0, heuristic: float = 0.0):
        self.state: T = state
        self.parent: Optional[Node] = parent
        self.cost: float = cost
        self.heuristic = heuristic

    def __lt__(self, other: 'Node') -> bool:
        return (self.cost + self.heuristic) < (other.cost + other.heuristic)


def node_to_path(node: Node[T]) -> List[T]:
    path: List[T] = [node.state]
    while node.parent is not None:
        node = node.parent
        path.append(node.state)
    path.reverse()
    return path


class Queue(Generic[T]):
    def __init__(self):
        self._container: Deque[T] = deque()

    @property
    def empty(self):
        return not self._container

    def pop(self):
        return self._container.popleft()

    def push(self, item: T):
        self._container.append(item)

    def __repr__(self):
        return repr(self._container)


@Function
def dfs(initial: T, goal_test: Callable[[T], bool], successors: Callable[[T], List[T]]) -> Optional[Node[T]]:
    """
    depth-first-search
    :param initial:
    :param goal_test:
    :param successors:
    :return:
    """
    frontier: Stack[Node[T]] = Stack()
    frontier.push(Node(initial, None))

    explored: Set[T] = {initial}

    while not frontier.empty:
        current_node: Node[T] = frontier.pop()
        current_state: T = current_node.state
        if goal_test(current_state):
            return current_node
        for child in successors(current_state):
            if child in explored:
                continue
            explored.add(child)
            frontier.push(Node(child, current_node))
    return None


@Function
def bfs(initial: T, goal_test: Callable[[T], bool], successors: Callable[[T], List[T]]) -> Optional[Node[T]]:
    """
    breadth-first-search
    :param initial:
    :param goal_test:
    :param successors:
    :return:
    """
    frontier: Queue[Node[T]] = Queue()
    frontier.push(Node(initial, None))
    explored: Set[T] = {initial}
    while not frontier.empty:
        current_node: Node[T] = frontier.pop()
        current_state: T = current_node.state
        if goal_test(current_state):
            return current_node
        for child in successors(current_state):
            if child in explored:
                continue
            explored.add(child)
            frontier.push(Node(child, current_node))
    return None


class PriorityQueue(Generic[T]):
    def __init__(self):
        self.__container: List[T] = []

    @property
    def empty(self):
        return not self.__container

    def pop(self):
        return heapq.heappop(self.__container)

    def push(self, item: T):
        heapq.heappush(self.__container, item)

    def __repr__(self):
        return repr(self.__container)


def euclidean_distance(goal: MazeLocation):
    def distance(ml: MazeLocation) -> float:
        xdist = ml.column - goal.column
        ydist = ml.row - goal.row
        return sqrt(xdist * xdist + ydist * ydist)

    return distance


def manhattan_distance(goal: MazeLocation) -> Callable[[MazeLocation], float]:
    def distance(ml: MazeLocation):
        xdist = abs(ml.column - goal.column)
        ydist = abs(ml.row - goal.row)
        return xdist + ydist

    return distance


@Function
def astar(initial: T, goal_test: Callable[[T], bool], successors: Callable[[T], List[T]],
          heuristic: Callable[[T], float]) -> Optional[Node[T]]:
    frontiers = PriorityQueue()
    frontiers.push(Node(initial, None, 0.0, heuristic(initial)))
    explored: Dict[T, float] = {initial: 0.0}
    while not frontiers.empty:
        current_node = frontiers.pop()
        current_state = current_node.state
        if goal_test(current_state):
            return current_node
        for child in successors(current_state):
            new_cost = current_node.cost + 1
            if child not in explored or explored[child] > new_cost:
                explored[child] = new_cost
                frontiers.push(Node(child, current_node, new_cost, heuristic(child)))


if __name__ == '__main__':
    st1 = {}
    st2 = {}
    st3 = {}
    zoom = 10 * 10
    for i in range(10, zoom + 1, 10):
        size = (i, i)
        m = Maze(rows=size[0], columns=size[1], goal=MazeLocation(size[0] - 1, size[1] - 1))
        distance = manhattan_distance(m.goal)
        solution = astar(m.start, m.goal_test, m.successors, distance)
        solution2 = bfs(m.start, m.goal_test, m.successors)
        solution3 = dfs(m.start, m.goal_test, m.successors)
        # print(solution.ret, solution2.ret)
        if solution2.ret is None or solution.ret is None or solution3.ret is None:
            print('No solution found ')
            # st1[i] = solution.spend_time
            # st2[i] = solution2.spend_time
            continue
        else:
            st1[i] = solution.spend_time
            st2[i] = solution2.spend_time
            st3[i] = solution3.spend_time
            # path = node_to_path(solution.ret)
            # m.mark(path)
            # print(m)
            # print(i, 'a star', solution.spend_time)
            # print(i, 'bfs   ', solution2.spend_time)
    import matplotlib.pyplot as plt

    plt.plot(st1.keys(), st1.values(), color='red', label='a_star')
    plt.plot(st2.keys(), st2.values(), color='blue', label='bfs')
    plt.plot(st3.keys(), st3.values(), color='green', label='dfs')
    plt.show()
    plt.savefig('time.jpg')
